import numpy as np
import matplotlib.pyplot as plt


class KMeans:
    def __init__(self, n_clusters=3, max_iter=100):
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.centroids = None
        self.labels_ = None

    def _distance(self, x1, x2):
        return np.sqrt(np.sum((x1 - x2) ** 2, axis=1))

    def _initialize_centroids(self, X):
        self.centroids = X[np.random.choice(len(X), self.n_clusters, replace=False)]

    def _assign(self, X):
        return np.array([np.argmin(self._distance(self.centroids, x)) for x in X])

    def _update_centroids(self, X, labels):
        return np.array([X[labels == i].mean(axis=0) for i in range(self.n_clusters)])

    def fit(self, X):
        self._initialize_centroids(X)
        for _ in range(self.max_iter):
            self.labels_ = self._assign(X)
            old_centroids = self.centroids.copy()
            self.centroids = self._update_centroids(X, self.labels_)
            if np.allclose(old_centroids, self.centroids, atol=1e-4):
                break
        return self


# 可直接使用（无需修改路径，基于生成/内置数据集）
if __name__ == "__main__":
    # 生成模拟数据
    np.random.seed(42)
    X = np.vstack([
        np.random.normal([2, 2], 0.5, (100, 2)),
        np.random.normal([5, 5], 0.5, (100, 2)),
        np.random.normal([8, 2], 0.5, (100, 2))
    ])

    kmeans = KMeans(n_clusters=3)
    kmeans.fit(X)
    print(f"聚类中心：\n{kmeans.centroids}")

    # 可视化
    plt.scatter(X[:, 0], X[:, 1], c=kmeans.labels_, cmap='viridis')
    plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1], c='red', marker='x')
    plt.show()